import pickle
import time
import itertools
import numpy as np
import random
import os

states_save_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'data/states.pkl')

board_positions = 9         # 棋盘的位置总数（从左到右，从上到下）
state_row_amount = 3        # 棋盘行数
state_column_amount = 3     # 棋盘列数
piece_state_amount = 3      # 棋盘上一个位置的可能状态数量 (0: 空，1:白子，2:黑子)


def checkListAllSameElements(l : list):
    """
    检查一个list的所有元素是否相等
    return bool
    """
    if len(l) < 1:
        return False
    
    old_e = None
    for e in l:
        if old_e is None:
            old_e = e
        else:
            if old_e != e:
                return False
            else:
                old_e = e

    return True

def checkStateIsWin(state):
    """
        检查棋盘的是否已赢（三子连成一线）
        返回值：（-1:没有赢家，只剩最后一个空位， 0:没有赢家，1:白子赢，2:黑子赢, 3:平局）
    """
    mat = np.array(state)
    mat = mat.reshape((state_row_amount, state_column_amount))
    # print(mat)

    # 检查棋盘的每行的棋子是否都相同
    for i in range(0, state_row_amount):
        l = mat[i]
        if checkListAllSameElements(l):
            # print('检查棋盘的每行的棋子是否都相同', l[0])
            if l[0] > 0:
                return l[0]
        
    # 检查棋盘的每列的棋子是否都相同
    for i in range(0, state_row_amount):
        l = mat[:, i]
        if checkListAllSameElements(l):
            # print('检查棋盘的每列的棋子是否都相同', l[0])
            if l[0] > 0:
                return l[0]

    # 检查对角线的棋子是否都相同
    l = np.diagonal(mat)
    if checkListAllSameElements(l):
        # print('检查对角线的棋子是否都相同', l[0])
        if l[0] > 0:
            return l[0]

    # 检查反向的对角线的棋子是否都相同
    l = np.diagonal(np.fliplr(mat))
    if checkListAllSameElements(l):
        # print('检查反向的对角线的棋子是否都相同', l[0])
        if l[0] > 0:
            return l[0]
    
    # 检查是否平局（没有空位置）
    if np.all(np.greater(state, [0] * 9)):
        return 3
    
    # 只剩一个空位
    if np.count_nonzero(np.array(state)) >= board_positions - 1:
        return -1
    
    return 0

def generateProbablity(state: tuple):
    
    zero_amount = state.count(0) # 空位置的个数
    
    # 平局
    if zero_amount == 0:
        return {}
    
    # 只剩下一个空位
    # if zero_amount == 0:
    #     return {}
    
    init_total = 100
    prob = int(init_total / zero_amount)

    d = {}

    for key, piece in enumerate(state):
        if piece == 0:
            d[key] = prob
        else:
            d[key] = 0

    return d

def generateInitStates():
    """
    生成初始 States,并保存到文件
    """

    states = {}

    # 生成所有的可能棋盘状态，笛卡尔积
    states_list = itertools.product(range(0, piece_state_amount), repeat=board_positions)
    
    # print(list(states_list))
    for state in list(states_list):
        isWin = checkStateIsWin(state)
        if 0 == isWin:
            states["".join(map(str, state))] = generateProbablity(state)

    setStates(states)

    return states;

def getStates():
    """
    从文件中读取 States
    """
    with open(states_save_file, 'rb') as f:
        return pickle.load(f)
    
def setStates(states):
    with open(states_save_file, 'wb') as f:
        pickle.dump(states, f)

# based on probability, to random select which position to put piece
def computeStateProb(states, stateName: str):
    state = list(map(int, list(stateName)))
    
    # 只剩一个空位的情况
    if np.count_nonzero(np.array(state)) >= board_positions - 1:
        for idx, piece in enumerate(state):
            if piece == 0:
                return idx

    
    probDict = states.get(stateName, {})
    if len(probDict) < 1:
        raise Exception('Can not find state board prob')

    total = 0
    for position in probDict.keys():
        total += probDict[position]


    ri = random.randint(1, total)
    # print('probDict:', probDict)
    # print('random:', ri)

    total = 0
    for position in probDict.keys():
        total += probDict[position]
        if ri <= total:
            return position

def changeStates(states: dict, machine_positions: dict, isWin: int):
    for stateName in machine_positions.keys():
        position = machine_positions[stateName]
        if stateName not in states:
            continue

        prob = states[stateName]
        print('prob change:', stateName, " from: ", prob)
        if isWin == 1:
            prob[position] -= 1
        elif isWin == 2:
            prob[position] += 2
        elif isWin == 3:
            prob[position] += 1
        
        if prob[position] <= 0:
            prob[position] = 1
        
        print('prob change:', stateName, " to  : ", prob)

        states[stateName] = prob

    setStates(states)

    return states


def listToStr(l: list):
    return "".join(map(str, l))

# get human input position
def getInputPosition(state: list):
    position = 0
    while True:
        print("Please input your position (from 1 to 9):")
        position = int(input())
        if position < 1 or position > 9:
            continue

        if state[position - 1] > 0:
            print("Error: this position is not empty!")
        else:
            break

    return position - 1

def printBoardState(state: list):
    mat = np.array(state)
    mat = mat.reshape((state_row_amount, state_column_amount))
    
    print("----------------------")
    for i in range(0, state_row_amount):
        print(mat[i])
    print("----------------------")


def printFirstHundredState(states):
    i = 0
    for key in states.keys():
        print(key, states[key])

        i += 1
        if i > 5:
            break
    # print(len(states))
    # exit();


if __name__ == "__main__":
    # states = generateInitStates()
    states = getStates()

    while True:
        state = [0] * 9 # blank board state
        whoesTurn = 2 #该谁下子，黑子先下; machine is black, human is white

        machinePositions = {};
        print(">>>>>>>>> START <<<<<<<<<<")
        print("Machine plays black(2), You plays white(1)")

        while True:
            # printState(state)

            # check whether there is a winner
            isWin = checkStateIsWin(state)
            if isWin > 0:
                if isWin == 1:
                    print("Game over, you wins!")
                elif isWin == 2:
                    print("Game over, Machine wins!")
                elif isWin == 3:
                    print("Game over, no winner!")
                
                states = changeStates(states, machinePositions, isWin)
                break
            else:
                # black‘s turn (which is machine)
                if whoesTurn == 2:
                    # compute position by probability
                    stateName = listToStr(state)
                    position = computeStateProb(states, stateName)
                    state[position] = whoesTurn
                    whoesTurn = 1
                    # record machine's positions
                    machinePositions[stateName] = position

                # white's turn (which is human)
                else:   
                    position = getInputPosition(state)
                    state[position] = whoesTurn
                    whoesTurn = 2
            printFirstHundredState(states)
            printBoardState(state)
